# -*- coding: utf-8 -*-  
'''
字典

@author: luoyi
Created on 2021年3月18日
'''
import pickle
import os

import utils.conf as conf


#    词<->id
word_to_index = {}
index_to_word = {}

#    追加一个可迭代的列表进dict
def append_to_dict(key_iter):
    global word_to_index
    crt_len = len(word_to_index)
    word_to_index.update({k : i + crt_len for i, k in enumerate(key_iter)})
    pass
#    追加一个可迭代的列表进dict，其意义与给定列表相同
def append_to_dict_synonym(keys, synonyms):
    global word_to_index
    word_to_index.update({k : word_to_index[s] for k, s in zip(keys, synonyms)})
    pass


#    词 <-> 索引    
# word_to_index = {
#                     #    初始文案: <PAD>填充字符 <UNK>未知字符 <GO>开始字符 <EOS>结束字符
#                     '[PAD]':0, '[UNK]':1, '[GO]':2, '[EOS]':3, '[MASK]':4,            
#                     #    追加0 ~ 9 10个数字
#                     '0':4, '1':5, '2':6, '3':7, '4':8, '5':9, '6':10, '7':11, '8':12, '9':13,
#                     #    追加a~z A~Z
#                     'a':14, 'b':15, 'c':16, 'd':17, 'e':18, 'f':19, 'g':20, 'h':21, 'i':22, 'g':23, 'k':24, 'l':25, 'm':26, 'n':27, 'o':28, 'p':29, 'q':30, 'r':31, 's':32, 't':33, 'u':34, 'v':35, 'w':36, 'x':37, 'y':38, 'z':39,
#                     'A':40, 'B':41, 'C':42, 'D':43, 'E':44, 'F':45, 'G':46, 'H':47, 'I':48, 'G':49, 'K':50, 'L':51, 'M':52, 'N':53, 'O':54, 'P':55, 'Q':56, 'R':57, 'S':58, 'T':59, 'U':60, 'V':61, 'W':62, 'X':63, 'Y':64, 'Z':65,
#                     #    追加中文标点符号，英文标点符号
#                     '。':66, '，':67, '、':68, '：':69, '；':70, '‘':71, '“':72, '？':73, '！':74, '……':75, '￥':76, '~':77, '《':78, '》':79, '（':80, '）':81, '「':82, '」':83,
#                     '.':84, ',':85, '/':86, ':':87, ';':88, "'":89, '"':90, '?':91, '!':92, '^':93, '$':94, '·':95, '<':96, '>':97, '(':98, ')':99, '[':100, ']':101, '{':102, '}':103,
#                     #    追加数学运算符
#                     '+':104, '-':105, '*':106, '@':107, '·':108, '×':109, '÷':110, '%':111, "//":112, '∑':113, '√':114, '∫':115, '▽':116, 'π':117, '∏':118
#                  }
#    追加[PAD] [UNK] [CLS] [EOS] [MASK]
append_to_dict('[PAD] [UNK] [CLS] [EOS] [MASK]'.split(' '))
#    追加0~9 10个数字
append_to_dict('0123456789')
append_to_dict_synonym('０１２３４５６７８９', '0123456789')
#    追加a~z A~Z
append_to_dict('abcdefghijklmnopqrstuvwxyz')
append_to_dict('abcdefghijklmnopqrstuvwxyz'.upper())
#    追加中文/英文标点符号
append_to_dict('！？｡＂＃￥＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏ .')
append_to_dict('.,/:;"?!^$<>()[]{}')
#    追加数学运算符
append_to_dict('+-*×÷@·%∑√∫▽π∏')

#    idx -> word
index_to_word = {v:k for k,v in word_to_index.items()}

#    是自动填充字典
is_auto_append = True
#    开启自动填充
def auto_append_on():
    global is_auto_append
    is_auto_append = True
    pass
#    关闭自动填充
def auto_append_off():
    global is_auto_append
    is_auto_append = False
    pass


#    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
def word2idx(word):
    idx = word_to_index.get(word)
    #    如果有就直接返回
    if (idx is not None): return idx
    
    #    如果不允许扩充字典，直接返回[UNK]
    if (not is_auto_append): 
        return word_to_index['[UNK]']
    
    #    否则扩充字典
    dict_size = len(word_to_index)
    word_to_index[word] = dict_size
    index_to_word[dict_size] = word
    
    return dict_size
#    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
def word2idx_slist(word_list):
    idx_list = [word2idx(word) for word in word_list]
    return idx_list

#    根据词的index取词
def id2word(word_id):
    return index_to_word.get(word_id)
#    根据词的index取词
def id2word_idlist(word_ids):
    id_list = [index_to_word.get(word_id) for word_id in word_ids]
    return id_list


#    字典长度
def dict_size():
    return len(word_to_index)


#    加载外部字典
def load_dict_from_pkl(word2dict_path=conf.TEXT.get_dictionaries_word2dict_pkl_path(),
                     dict2word_path=conf.TEXT.get_dictionaries_dict2word_pkl_path()):
    global word_to_index
    global index_to_word
    with open(word2dict_path, mode='rb') as file_word2dict:
        word_to_index.update(pickle.load(file_word2dict))
    with open(dict2word_path, mode='rb') as file_dict2word:
        index_to_word.update(pickle.load(file_dict2word))
    pass
#    保存字典pkl
def save_dict_to_pkl(word2dict_path=conf.TEXT.get_dictionaries_word2dict_pkl_path(),
                     dict2word_path=conf.TEXT.get_dictionaries_dict2word_pkl_path()):
    #    建立目录
    conf.mkfiledir_ifnot_exises(word2dict_path)
    conf.mkfiledir_ifnot_exises(dict2word_path)
    #    删除同名文件
    if (os.path.exists(word2dict_path)): os.remove(word2dict_path)
    if (os.path.exists(dict2word_path)): os.remove(dict2word_path)
    with open(word2dict_path, mode='wb') as fw:
        pickle.dump(word_to_index, fw)
    with open(dict2word_path, mode='wb') as fw:
        pickle.dump(index_to_word, fw)
    pass


#    加载外部字典
def load_dict_from_chchars(chchars_path=conf.TEXT.get_dictionaries_chchars_path()):
    fr = open(file=chchars_path, mode='r', encoding='utf-8')
    i = 0
    while True:
        l = fr.readline()
        if (not l): break
        if (len(l) == 0): continue
        
        word = l.split(':')[0]
        word2idx(word)
        i += 1
        pass
    fr.close()
    print('读到汉字:{}'.format(i))
    pass


    